﻿using System;
using System.Collections.Generic;
using System.Drawing;
using System.Linq;
using System.Runtime.InteropServices;
using System.Security;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Forms;
using FreeImageAPI;
using OpenTK;
using OpenTK.Graphics.OpenGL;
using OpenTK.Input;
using OptixDotNet;
using OptixDotNet.Math;
using OptixDotNet.Utils;
using OptixDotNetUtils;
using Vector3 = OpenTK.Vector3;
using Vector4 = OpenTK.Vector4;

namespace RayDen.Optix.Engine.Components
{
    public class BaseWindow : GameWindow {

        public class FrameManager {
            public TimeSpan TimeForFrame = new TimeSpan(0,0,3);
            public BaseWindow owner;
            public int framesCount;
            public int currentFrame;
            public string FrameNameTemplate = "anim_{0}.hdr";

            OptixDotNet.Math.Vector3 startP, startD;
            OptixDotNet.Math.Vector3 endP, endD;

            private float plength, dlength;
            private bool isRuning;
            private DateTime? startTime;

            public FrameManager(OptixDotNet.Math.Vector3 sp, OptixDotNet.Math.Vector3 sd, OptixDotNet.Math.Vector3 ep, OptixDotNet.Math.Vector3 ed, BaseWindow owner) {
                this.owner = owner;
                currentFrame = 0;
                framesCount = 30;
                plength = (ep - sp).Length() / (float)framesCount;
                dlength = (ed - sd).Length() / (float)framesCount;
                startP = sp;
                endP = ep;
                startD = sd;
                endD = ed;
            }

            public static float Lerp(float black, float red, float step) {
                return step * black + (1 - step) * red;
            }

            public void GetCurrentValues(out OptixDotNet.Math.Vector3 pos, out OptixDotNet.Math.Vector3 dir) {
                float pstep = currentFrame * plength;
                float dstep = currentFrame * dlength;

                pos = startP +
                    new OptixDotNet.Math.Vector3(
                        Lerp(endP.X, startP.X, pstep),
                        Lerp(endP.Y, startP.Y, pstep),
                        Lerp(endP.Z, startP.Z, pstep)
                        );
                dir = startD +
                    new OptixDotNet.Math.Vector3(
                        Lerp(endD.X, startD.X, dstep),
                        Lerp(endD.Y, startD.Y, dstep),
                        Lerp(endD.Z, startD.Z, dstep)
                        );
            }

            public void UpdateCamera(ICamera cam) {
                float pstep = currentFrame * plength;
                float dstep = currentFrame * dlength;
                cam.Dolly(pstep);
            }


            internal void Reset() {
                startTime = null;
                isRuning = false;
                currentFrame = 0;
            }

            public void StartRendering() {
                if (isRuning)
                    return;
                isRuning = true;
                startTime = DateTime.UtcNow;
                currentFrame = 0;
            }

            public bool IsFrameReady() {
                return isRuning && ((DateTime.UtcNow - startTime.GetValueOrDefault()) > TimeForFrame);
            }

            public void UpdateFrame() {
                if (currentFrame >= framesCount || !isRuning)
                    return;
                //owner.SavePBOhdr(string.Format(this.FrameNameTemplate, currentFrame));
                currentFrame++;
                startTime = DateTime.UtcNow;
            }
        }

        [SuppressUnmanagedCodeSecurity]
        [DllImport("user32.dll", CharSet = CharSet.Auto)]
        [return: MarshalAs(UnmanagedType.Bool)]
        public static extern bool PeekMessage(out Message msg, IntPtr hWnd, uint messageFilterMin, uint messageFilterMax, uint flags);


        public ICamera Camera = null;
        protected string ImageSaveDir;
        public Context Context = null;
        public OptixDotNet.Buffer OutputBuffer = null;

        public Vector3 BackBufferColor;


        public bool UsePBO = true;
        public bool UseSRGB = false;
        public bool DrawUI = true;
        #region Private Fields

        //stats
        public int RaysTracedPerFrame;
        private int mNumFrames = 0;
        private int mNumRaysTraced = 0;
        private float mTimeElapsed = 0.0f;
        private float mTimeDelta = 0.0f;
        private string mText = "";

        //ogl
        private uint mPBOTexID = 0;

        //mouse
        private MouseButtonEventArgs mPrevMouse;

        protected FrameManager frameManager;
        #endregion

        public BaseWindow()
            : base(1280, 720) {

            Time.InitTime();
            Keyboard.KeyDown += new EventHandler<OpenTK.Input.KeyboardKeyEventArgs>(Keyboard_KeyDown);
            Mouse.ButtonDown += new EventHandler<MouseButtonEventArgs>(Mouse_Move);
            mPrevMouse = new MouseButtonEventArgs();
        }


        protected virtual void Mouse_Move(object sender, MouseButtonEventArgs e) {
            if (Camera == null)
                return;

            float deltaX = (e.X - mPrevMouse.X) / (float)Width;
            float deltaY = (e.Y - mPrevMouse.Y) / (float)Height;

            deltaX *= mTimeDelta;
            deltaY *= mTimeDelta;

            float deltaT = deltaX;
            if (Math.Abs(deltaX) < Math.Abs(deltaY))
                deltaT = deltaY;

            switch (e.Button) {
                case MouseButton.Left:
                    Camera.Rotate(-deltaX, -deltaY);
                    CameraUpdate();
                    break;
                case MouseButton.Right:
                    Camera.Dolly(-deltaT);
                    CameraUpdate();
                    break;
                case MouseButton.Middle:
                    Camera.Pan(-deltaX, -deltaY);
                    CameraUpdate();
                    break;
                default:
                    break;
            }

            mPrevMouse = e;
        }

        protected virtual void Keyboard_KeyUp(object sender, KeyboardKeyEventArgs e) {

        }

        protected virtual void Keyboard_KeyDown(object sender, OpenTK.Input.KeyboardKeyEventArgs e) {
            if (Camera == null)
                return;

            float step = mTimeDelta *0.1f;

            float angle = 10.0f * (float)Math.PI / 180f;
            switch (e.Key) {
                case Key.Up:
                    Camera.Dolly(step);
                    CameraUpdate();

                    break;
                case Key.Down:
                    Camera.Dolly(-step);
                    CameraUpdate();

                    break;
                case Key.Left:
                    Camera.Pan(step, 0.0f);
                    CameraUpdate();

                    break;
                case Key.Right:
                    Camera.Pan(-step, 0.0f);
                    CameraUpdate();

                    break;


                case Key.Comma:
                    Camera.Pan(0.0f, step);
                    CameraUpdate();

                    break;
                case Key.Period:
                    Camera.Pan(0.0f, -step);
                    CameraUpdate();

                    break;

                case Key.F2:
                    this.SetCamera(new BoundingBox());
                    break;

                case Key.F5:
                    var ticks = DateTime.UtcNow.Ticks;
                    string bitmapName = string.Format("{1}\\frame_{0}.png", ticks, ImageSaveDir);
                    string HDRName = string.Format("{1}\\frame_{0}.hdr", ticks, ImageSaveDir);

                    SavePBO(bitmapName);
                    SavePBOhdr(HDRName);
                    return;
                    break;
                case Key.Escape:
                    Close();
                    break;
                case Key.Delete:
                    this.Camera.Rotate(0f, angle);
                    CameraUpdate();

                    break;
                case Key.PageDown:
                    this.Camera.Rotate(0f, -angle);
                    CameraUpdate();

                    break;
                case Key.End:
                    this.Camera.Rotate(angle, 0f);
                    CameraUpdate();

                    break;
                case Key.Home:
                    this.Camera.Rotate(-angle, 0f);
                    CameraUpdate();
                    break;
                default:
                    break;
            }

        }

        protected override void OnLoad(EventArgs e) {
            base.OnLoad(e);
            GL.Enable(EnableCap.FramebufferSrgb);
            //Application.Idle += OnIdle;
            Initialize();
        }

        protected virtual void Initialize() { }
        protected virtual void RayTrace() { }
        protected virtual void SetCamera(BoundingBox box) { }

        private string stats;
        public virtual string GetStatistics() {
            return this.stats;
        }
        private bool AppStillIdle() {
            Message msg;
            return !PeekMessage(out msg, IntPtr.Zero, 0, 0, 0);
        }

        public void OnIdle(object sender, EventArgs e) {
            while (AppStillIdle()) {
                start = Time.GetTimeInSecsDouble();
                RayTrace();
                stop = Time.GetTimeInSecsDouble();

                mTimeElapsed += (float)(stop - start);
                mNumFrames++;
                mNumRaysTraced += RaysTracedPerFrame;

                if (mTimeElapsed >= 0.25f) {
                    float raysPerSec = mNumRaysTraced / mTimeElapsed;
                    raysPerSec /= 1000.0f * 1000.0f;

                    stats = string.Format("ray-trace: {0:.00} ms : {1:.00} MRays/sec {2:.000} frames per pixel", mTimeElapsed / mNumFrames * 1000.0f, raysPerSec, mNumFrames / (double)raysPerSec);
                    mTimeElapsed = 0.0f;
                    mNumFrames = 0;
                    mNumRaysTraced = 0;
                }

                // Initialize state
                GL.MatrixMode(MatrixMode.Projection);
                GL.LoadIdentity();
                GL.Ortho(0, 1, 0, 1, -1, 1);
                GL.MatrixMode(MatrixMode.Modelview);
                GL.LoadIdentity();
                GL.Viewport(0, 0, Width, Height);

                GL.ClearColor(BackBufferColor.X, BackBufferColor.Y, BackBufferColor.Z, 1);
                GL.Clear(ClearBufferMask.ColorBufferBit);

                Display();

                DrawText(15, 15, GetStatistics());
                this.Title = stats;
                System.Threading.Thread.Sleep(10);
                //SwapBuffers();

                stop = Time.GetTimeInSecsDouble();
                mTimeDelta = (float)(stop - start);
            }
        }

        private double start, stop;
        protected override void OnUpdateFrame(FrameEventArgs e) {
            base.OnUpdateFrame(e);
            start = Time.GetTimeInSecsDouble();
            try
            {
                RayTrace();
#if MULTIFRAMES
            if (frameManager.IsFrameReady()) {
                frameManager.UpdateFrame();
                OptixDotNet.Math.Vector3 cp, cl;
                //frameManager.UpdateCamera(Camera);
                frameManager.GetCurrentValues(out cp, out cl);
                Camera.Position = cp;
                Camera.Look = cl;
                CameraUpdate();
            }
#endif
                stop = Time.GetTimeInSecsDouble();

                mTimeElapsed += (float) (stop - start);
                mNumFrames++;
                mNumRaysTraced += RaysTracedPerFrame;

                if (mTimeElapsed >= 0.25f)
                {
                    float raysPerSec = mNumRaysTraced/mTimeElapsed;
                    raysPerSec /= 1000.0f*1000.0f;

                    stats = string.Format("ray-trace: {0:.00} ms : {1:.00} MRays/sec {2:.000} frames per pixel",
                        mTimeElapsed/mNumFrames*1000.0f, raysPerSec, mNumFrames/(double) raysPerSec);
                    mTimeElapsed = 0.0f;
                    mNumFrames = 0;
                    mNumRaysTraced = 0;
                }
                mTimeDelta = (float) (stop - start);
            }
            catch (Exception ex)
            {
                this.Exit();
            }
        }
        protected override void OnRenderFrame(FrameEventArgs e) {
            base.OnRenderFrame(e);

            // Initialize state
            GL.MatrixMode(MatrixMode.Projection);
            GL.LoadIdentity();
            GL.Ortho(0, 1, 0, 1, -1, 1);
            GL.MatrixMode(MatrixMode.Modelview);
            GL.LoadIdentity();
            GL.Viewport(0, 0, Width, Height);

            GL.ClearColor(BackBufferColor.X, BackBufferColor.Y, BackBufferColor.Z, 1);
            GL.Clear(ClearBufferMask.ColorBufferBit);

            Display();

            this.Title = this.GetStatistics();
            DrawText(15, 15, GetStatistics());

            System.Threading.Thread.Sleep(10);

        }


        protected void DrawPBO() {
            OGLBuffer buffer = OutputBuffer as OGLBuffer;

            if (buffer == null)
                return;

            if (mPBOTexID == 0) {
                GL.GenTextures(1, out mPBOTexID);
                GL.BindTexture(TextureTarget.Texture2D, mPBOTexID);

                // Change these to GL_LINEAR for super- or sub-sampling
                GL.TexParameter(TextureTarget.Texture2D, TextureParameterName.TextureMagFilter, (int)TextureMagFilter.Nearest);
                GL.TexParameter(TextureTarget.Texture2D, TextureParameterName.TextureMinFilter, (int)TextureMinFilter.Nearest);

                // GL_CLAMP_TO_EDGE for linear filtering, not relevant for nearest.
                GL.TexParameter(TextureTarget.Texture2D, TextureParameterName.TextureWrapS, (int)TextureWrapMode.ClampToEdge);
                GL.TexParameter(TextureTarget.Texture2D, TextureParameterName.TextureWrapT, (int)TextureWrapMode.ClampToEdge);

                GL.BindTexture(TextureTarget.Texture2D, 0);
            }

            GL.BindTexture(TextureTarget.Texture2D, mPBOTexID);

            // send pbo to texture
            GL.BindBuffer(BufferTarget.PixelUnpackBuffer, buffer.Resource);

            var elementSize = buffer.ElementSize;
            if ((elementSize % 8) == 0) GL.PixelStore(PixelStoreParameter.UnpackAlignment, 8);
            else if ((elementSize % 4) == 0) GL.PixelStore(PixelStoreParameter.UnpackAlignment, 4);
            else if ((elementSize % 2) == 0) GL.PixelStore(PixelStoreParameter.UnpackAlignment, 2);
            else GL.PixelStore(PixelStoreParameter.UnpackAlignment, 1);

            switch (buffer.Format) {
                case Format.UByte4:
                    GL.TexImage2D(TextureTarget.Texture2D, 0, PixelInternalFormat.Rgba8, Width, Height, 0, PixelFormat.Bgra, PixelType.UnsignedByte, IntPtr.Zero);
                    break;
                case Format.Float4:
                    GL.TexImage2D(TextureTarget.Texture2D, 0, PixelInternalFormat.Rgba32f, Width, Height, 0, PixelFormat.Rgba, PixelType.Float, IntPtr.Zero);
                    break;
                case Format.Float3:
                    GL.TexImage2D(TextureTarget.Texture2D, 0, PixelInternalFormat.Rgba32f, Width, Height, 0, PixelFormat.Rgb, PixelType.Float, IntPtr.Zero);
                    break;
                case Format.Float:
                    GL.TexImage2D(TextureTarget.Texture2D, 0, PixelInternalFormat.Luminance, Width, Height, 0, PixelFormat.Luminance, PixelType.Float, IntPtr.Zero);
                    break;
                default:
                    throw new Exception("Unsupported format for PBO");
                    break;
            }

            //GL.TexImage2D( TextureTarget.Texture2D, 0, Gl.GL_RGBA8, Width, Height, 0, Gl.GL_BGRA, Gl.GL_UNSIGNED_BYTE, IntPtr.Zero );

            GL.BindBuffer(BufferTarget.PixelUnpackBuffer, 0);

            GL.Enable(EnableCap.Texture2D);

            // Initialize offsets to pixel center sampling.

            float u = 0.5f / Width;
            float v = 0.5f / Height;

            GL.Begin(BeginMode.Quads);
            GL.TexCoord2(u, v);
            GL.Vertex2(0.0f, 0.0f);
            GL.TexCoord2(1.0f, v);
            GL.Vertex2(1.0f, 0.0f);
            GL.TexCoord2(1.0f - u, 1.0f - v);
            GL.Vertex2(1.0f, 1.0f);
            GL.TexCoord2(u, 1.0f - v);
            GL.Vertex2(0.0f, 1.0f);
            GL.End();

            GL.Disable(EnableCap.Texture2D);
        }
        protected virtual void Display() {
            if (UsePBO) {
                DrawPBO();
            }
            else {
                BufferStream stream = OutputBuffer.Map();

                PixelType type = 0;
                PixelFormat format = 0;
                switch (OutputBuffer.Format) {
                    case Format.UByte4:
                        format = PixelFormat.Bgra;
                        type = PixelType.UnsignedByte;
                        break;
                    case Format.Float4:
                        format = PixelFormat.Rgba;
                        type = PixelType.Float;
                        break;
                    case Format.Float3:
                        format = PixelFormat.Rgb;
                        type = PixelType.Float;
                        break;
                    case Format.Float:
                        format = PixelFormat.Luminance;
                        type = PixelType.Float;
                        break;
                    default:
                        throw new Exception("Unsupported format for glDrawPixels");
                        break;
                }

                GL.DrawPixels(Width, Height, format, type, stream.DataPointer);

                OutputBuffer.Unmap();
            }
            SwapBuffers();
        }

        protected override void OnClosed(EventArgs e) {
            base.OnClosed(e);
            Application.Idle -= OnIdle;
            //Context.Destroy();
            //Context = null;
        }
        protected virtual void DrawText(int x, int y, string text) {
            if (!DrawUI)
                return;
            /*
            GL.MatrixMode(Gl.GL_PROJECTION);
            GL.LoadIdentity();
            GL.Ortho(0, Width, 0, Height, -1, +1);

            // Save state
            GL.PushAttrib(Gl.GL_CURRENT_BIT | Gl.GL_ENABLE_BIT);

            GL.Disable(TextureTarget.Texture2D);
            GL.Disable(Gl.GL_LIGHTING);
            GL.Disable(Gl.GL_DEPTH_TEST);

            GL.Color3f(.1f, .1f, .1f); // drop shadow color

            // Shift shadow one pixel to the lower right.
            GL.RasterPos2i(x + 1, y - 1);
            foreach (char c in text)
                Glut.glutBitmapCharacter(Glut.GLUT_BITMAP_8_BY_13, (int)c);

            GL.Color3f(1.0f, 1.0f, 1.0f);        // main text
            GL.RasterPos2i(x, y);
            foreach (char c in text)
                Glut.glutBitmapCharacter(Glut.GLUT_BITMAP_8_BY_13, (int)c);

            // Restore state
            GL.PopAttrib();
             * */
        }


        protected virtual void CameraUpdate() {

        }

        protected void CreateOutputBuffer(Format format) {
            if (UsePBO) {
                //Allocate first the memory for the gl buffer, then attach it to OptiX.
                uint pbo = CreatePBO(format);

                OGLBufferDesc glDesc = new OGLBufferDesc() {
                    Width = (uint)Width,
                    Height = (uint)Height,
                    Format = format,
                    Type = BufferType.Output,
                    Resource = pbo
                };
                OutputBuffer = new OGLBuffer(Context,  glDesc);
            }
            else {
                BufferDesc desc = new BufferDesc() {
                    Width = (uint)Width,
                    Height = (uint)Height,
                    Format = format,
                    Type = BufferType.Output
                };
                OutputBuffer = EngineContext.Instance.ResourceManager.CreateBuffer(Context, desc);
            }
        }

        protected uint CreatePBO(Format format) {
            uint pbo = 0;
            GL.GenBuffers(1, out pbo);
            GL.BindBuffer(BufferTarget.ArrayBuffer, pbo);

            int size = (int)Utils.GetFormatSize(format);
            GL.BufferData(BufferTarget.ArrayBuffer, new IntPtr(size * Width * Height), IntPtr.Zero, BufferUsageHint.StreamDraw);
            GL.BindBuffer(BufferTarget.ArrayBuffer, 0);

            return pbo;
        }
        protected void SavePBOhdr(string fileName) {
            OGLBuffer buffer = OutputBuffer as OGLBuffer;
            if (buffer == null)
                return;
            Vector4[] data = new Vector4[buffer.Width * buffer.Height];
            buffer.GetDataNoAlloc<Vector4>(data);
            using (var bmp = new FreeImageBitmap(this.Width, this.Height, FREE_IMAGE_TYPE.FIT_RGBF)) {
                for (var j = 0ul; j < buffer.Height; j++) {
                    var scanLine = bmp.GetScanline<FIRGBF>((int)buffer.Height - (int)j - 1);
                    for (var i = 0ul; i < buffer.Width; i++) {
                        var px = data[(int)i + (int)(buffer.Height - j - 1) * bmp.Width];
                        scanLine[(int)i] = new FIRGBF() { red = px.X, green = px.Y, blue = px.Z};
                    }
                }
                bmp.Save(fileName, FREE_IMAGE_FORMAT.FIF_EXR);
            }
        }
        public static float Clamp(float cosi, float p, float p_3) {
            if (cosi < p)
                return p;
            if (cosi > p_3)
                return p_3;
            return cosi;
        }

        public static int Clamp(int cosi, int p, int p_3) {
            int r = cosi;
            if (r < p)
                r = p;
            if (r > p_3)
                r = p_3;

            return r;

        }
        private static int GAMMA_TABLE_SIZE = 1024;
        float[] gammaTable;
        internal void InitGammaTable() {
            float x = 0f;
            gammaTable = new float[GAMMA_TABLE_SIZE];
            float dx = 1f / GAMMA_TABLE_SIZE;
            for (int i = 0; i < GAMMA_TABLE_SIZE; ++i, x += dx)
                gammaTable[i] = (float)Math.Pow(Clamp(x, 0f, 1f), 1f / 2.2f);
        }

        internal float Radiance2PixelFloat(float x) {
            // Very slow !
            //return powf(Clamp(x, 0.f, 1.f), 1.f / 2.2f);

            int index = Math.Min(
                Floor2UInt(GAMMA_TABLE_SIZE * Clamp(x, 0f, 1f)),
                    GAMMA_TABLE_SIZE - 1);
            return gammaTable[index];
        }

        public static int Floor2UInt(float val) {
            return val > 0 ? (int)(val) : 0;
        }

        protected void SavePBO(string fileName) {
            InitGammaTable();
            OGLBuffer buffer = OutputBuffer as OGLBuffer;

            if (buffer == null)
                return;
            Vector4[] data = new Vector4[buffer.Width * buffer.Height];
            buffer.GetDataNoAlloc<Vector4>(data);

            using (var bmp = new Bitmap(this.Width, this.Height)) {
                for (var j = 0ul; j < buffer.Height; j++)
                    for (var i = 0ul; i < buffer.Width; i++) {
                        bmp.SetPixel((int)i, (int)j, Vector4ToColor(data[(int)i + ((int)(buffer.Height - j - 1)) * bmp.Width]));
                    }
                bmp.Save(fileName);
            }

        }

        private Color Vector4ToColor(Vector4 c) {
            return Color.FromArgb((byte)(Radiance2PixelFloat(c.X) * 255.0f), (byte)(Radiance2PixelFloat(c.Y) * 255.0f), (byte)(Radiance2PixelFloat(c.Z) * 255.0f));
        }
    }
}
